import wordcloud
import base64
import numpy
from PIL import Image
import io
import sys


# 依赖安装: pip install wordcloud

# 参数text: 词云内容及频率组成的字符串 "内容1,频率1 内容2,频率2 内容3,频率3 ..."
# 参数mask_pic: 底图的地址
# 通过让某个词语在text中多次出现以保证其在词云图中占据较大比例
# 函数返回值为base64编码的字符串,可以通过在html中设置标签 <img src="data:image/png;base64,###"> 来展示词云图片(把###替换成函数返回的字符串)
def get_word_cloud_base64(text, mask_pic):
    # 设置底图
    mask = numpy.array(Image.open(mask_pic))
    # 字符串处理
    frequencies = {}
    text = text.split(" ")
    for i in text:
        arr = i.split(",")
        frequencies[arr[0]] = float(arr[1])
    # 使用wordcloud生成词云图
    pil_img = wordcloud.WordCloud(mask=mask, width=1100, height=670, background_color="white", font_path="msyh.ttc").generate_from_frequencies(frequencies).to_image()
    img = io.BytesIO()
    pil_img.save(img, "PNG")
    img.seek(0)
    img_base64 = base64.b64encode(img.getvalue()).decode()
    return img_base64

def get_word_cloud_base64_random(text, mask_pic):
    # 设置底图
    mask = numpy.array(Image.open(mask_pic))
    # 使用wordcloud生成词云图
    pil_img = wordcloud.WordCloud(mask=mask, width=1100, height=670, background_color="white", font_path="msyh.ttc").generate(text=text).to_image()
    img = io.BytesIO()
    pil_img.save(img, "PNG")
    img.seek(0)
    img_base64 = base64.b64encode(img.getvalue()).decode()
    return img_base64

# 参数text: 词云内容组成的字符串,由多个词语间隔空格组成
# 参数mask_pic: 底图的地址
# 传入字符串生成并保存词云图,保存的文件名为file_name,函数无返回值 例: save_word_cloud_pics("反复咳嗽 咳痰 喘息加重3天", "my_word_cloud.png")
def save_word_cloud_pics(text, mask_pic, file_name):
    # 设置底图
    mask = numpy.array(Image.open(mask_pic))
    wordcloud.WordCloud(mask=mask, width=1100, height=670, background_color="white", font_path="msyh.ttc").generate(text=text).to_file(file_name)
    
    
if __name__=="__main__":
    mask_pic = './ecas/src/main/resources/static/mask_square.png'
    if sys.argv[1] == 'frequencies':
        print(get_word_cloud_base64(sys.argv[2], mask_pic))
    else:
        print(get_word_cloud_base64_random(sys.argv[2], mask_pic))
    